import datetime
import math
import numpy as np
import torch

from torch.nn import MSELoss
from torch.optim import Adam
from os import path

from ARM_TGNCF.ARM_TGNCF_model import ARM_TGNCF
from ARM_TGNCF.Parse_ARM_TGNCF import ARM_args
from ARM_TGNCF.tool import initNetParams
from loaddata.load_data import load_ratings_data
from tool import get_UI_type, get_DataLoader, print_args_information


def begin():
    torch.autograd.set_detect_anomaly(True)

    SEED = ARM_args.seed
    torch.manual_seed(SEED)
    torch.cuda.manual_seed(SEED)
    np.random.seed(SEED)

    relation_uu, relation_ii, method = get_UI_type()

    train_event, test_event, ratings_num = load_ratings_data()

    user_num, item_num, dl_train, dl_test = get_DataLoader(train_event, test_event)

    if ARM_args.is_use_cuda:
        model = ARM_TGNCF(user_num, item_num, train_event, ARM_args.embed_size,
                          layers=[ARM_args.embed_size, ] * ARM_args.lysz).cuda()
    else:
        model = ARM_TGNCF(user_num, item_num, train_event, ARM_args.embed_size,
                          layers=[ARM_args.embed_size, ] * ARM_args.lysz, is_use_cuda=False)

    optim = Adam(model.parameters(), lr=ARM_args.lr, weight_decay=ARM_args.decay)

    loss_fn = MSELoss()

    info = print_args_information(method, user_num, item_num, ratings_num)

    title = str(ARM_args.ii_method) + '_' + str(ARM_args.dataset) + '_' + str(ARM_args.batch_size) + '_' + str(
        ARM_args.forecast_method) + '_' + str(ARM_args.embed_size)
    if relation_uu:
        title += '_u'
    if relation_ii:
        title += '_i'
    if not (relation_ii or relation_uu):
        title += '_None'
    if ARM_args.is_in_train:
        title += ('_' + str(ARM_args.is_in_train))

    result_file_path = path.join(path.dirname(__file__), 'result', title + '.txt')

    best_model_path = path.join(path.dirname(__file__), 'result', 'model_result', title + '_best')

    cur_model_path = best_model_path + '_cur'

    if path.exists(cur_model_path) and not ARM_args.model_is_rerun:
        checkpoint = torch.load(cur_model_path)
        model.load_state_dict(checkpoint['net'])
        optim.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch'] + 1
        min_mae = checkpoint['min_mae']
        max_recall = checkpoint['max_recall']
        max_accuracy = checkpoint['max_accuracy']

    else:
        min_mae = float('inf')
        max_recall = 0.0
        max_accuracy = 0.0
        start_epoch = 0

        # 保存试验提示信息
        result_file = open(result_file_path, 'w')
        result_file.write(info)
        result_file.close()

        # 初始化模型参数
        initNetParams(model)

    # 开始迭代
    for epoch in range(start_epoch, ARM_args.epoch + 1):
        begin_time = datetime.datetime.now()
        final_emb = model.getFeatureMat()
        events_len = 0  # 已经训练过的事件数
        for batch_idx, batch in enumerate(dl_train):
            if ARM_args.is_in_train:
                events_len += len(batch[1])
                events = train_event.head(events_len)
                model.LaplacianMat, model.ui_LaplacianMat = \
                    model.syn_build_LaplacianMat(events,
                                                 final_embed=final_emb,
                                                 uu_flg=relation_uu,
                                                 ii_flg=relation_ii)
            optim.zero_grad()
            if ARM_args.is_use_cuda:
                user_idx = batch[0].cuda()
                item_idx = batch[1].cuda()
                # rating_idx = torch.Tensor(batch[2]).float().cuda()
                rating_idx = batch[2].float().cuda()
            else:
                user_idx = batch[0]
                item_idx = batch[1]
                rating_idx = batch[2].float()
            prediction, user_emb, item_emb, final_emb = model(user_idx, item_idx)
            loss = loss_fn(rating_idx, prediction)
            for param in model.parameters():
                loss = loss + 0.0001 * param.norm(2) ** 2
            loss.backward()
            optim.step()

        # if i > int(args.epoch) * args.dk:
        if epoch >= ARM_args.test_start_epoch:
            info = 'epoch:{}'.format(epoch)

            model.LaplacianMat, model.ui_LaplacianMat = \
                model.syn_build_LaplacianMat(train_event,
                                             final_embed=final_emb,
                                             uu_flg=relation_uu,
                                             ii_flg=relation_ii)

            for data in dl_test:
                if ARM_args.is_use_cuda:
                    user_idx = data[0].cuda()
                    item_idx = data[1].cuda()
                    ratings_idx = data[2].float().cuda()
                else:
                    user_idx = data[0]
                    item_idx = data[1]
                    ratings_idx = data[2].float()
                prediction, user_emb, item_emb, final_emb_2 = model(user_idx, item_idx)
                # 保存模型
                state = {'net': model.state_dict(),
                         'optimizer': optim.state_dict(),
                         'epoch': epoch,
                         'min_mae': min_mae,
                         'max_recall': max_recall,
                         'max_accuracy': max_accuracy
                         }

                if ARM_args.forecast_method == 'inter':
                    TP = 0
                    FP = 0
                    FN = 0
                    TN = 0
                    predictions = list(prediction)
                    labels = list(data[2])
                    for index in range(len(predictions)):
                        pred = predictions[index]
                        label = labels[index]
                        if label == 1:
                            if pred > 0.5:
                                TP += 1
                            else:
                                FN += 1
                        else:
                            if pred > 0.5:
                                FP += 1
                            else:
                                TN += 1

                    assert len(prediction) == TP + FP + TN + FN
                    if TP + FP > 0:
                        recall = TP / (TP + FP)
                    else:
                        recall = 0.0
                    accuracy = (TP + TN) / (TP + FP + TN + FN)
                    info += '  recall:{}  accuracy:{} '.format(recall, accuracy)
                    info += '  正例:{}  负例:{} '.format(TP + FN, FP + TN)
                    info += '  time:{}'.format(datetime.datetime.now() - begin_time)
                    # 保存最优值的模型
                    if recall > max_recall:
                        max_recall = recall
                        state['max_recall'] = max_recall
                        state['max_accuracy'] = max_accuracy
                        torch.save(state, best_model_path)

                else:
                    mse = loss_fn(ratings_idx, prediction).cpu().detach().numpy()
                    mae = np.mean(np.abs(data[2].detach().numpy() - prediction.cpu().detach().numpy()))
                    rmse = math.sqrt(mse)
                    info += '  rmse:{}  mse:{}  mae:{}'.format(rmse, mse, mae)
                    info += '  time:{}'.format(datetime.datetime.now() - begin_time)
                    # 保存最优值的模型
                    if mae < min_mae:
                        min_mae = mae
                        state['min_mae'] = min_mae
                        torch.save(state, best_model_path)

                # 保存当前模型
                torch.save(state, cur_model_path)

                #  保存输出结果
                result_file = open(result_file_path, 'a')
                result_file.write(info + '\n')
                result_file.close()
                print(info)


if __name__ == '__main__':
    begin()
